/*
 * Copyright 2021 <copyright holder> <email>
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <iostream>
#include <fstream>
#include <string>
#include <eigen3/Eigen/Core>

#include <g2o/core/base_vertex.h>
#include <g2o/core/base_binary_edge.h>
#include <g2o/core/block_solver.h>
#include <g2o/core/optimization_algorithm_levenberg.h>
#include <g2o/solvers/eigen/linear_solver_eigen.h>

#include <sophus/se3.hpp>

using namespace std;
using namespace Eigen;
using Sophus::SE3d;
using Sophus::SO3d;

typedef Matrix<double,6,6> Matrix6d;

Matrix6d JRInv(const SE3d &e){
    Matrix6d J;
    J.block(0,0,3,3)=SO3d::hat(e.so3().log());
    J.block(0,3,3,3)=SO3d::hat(e.translation());
    J.block(3,0,3,3)=Matrix3d::Zero(3,3);
    J.block(3,3,3,3)=SO3d::hat(e.so3().log());
    //J=J*0.5+Matrix6d::Identity();
    J=Matrix6d::Identity();
    return J;
}

typedef Matrix<double,6,1>Vector6d;

class VertexSE3LieAlgebra: public g2o::BaseVertex<6,SE3d>{
public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    
    virtual bool read(istream &is) override{
        double data[7];
        for(int i=0;i<7;i++)
            is>>data[i];
        setEstimate(SE3d(
            Quaterniond(data[6],data[3],data[4],data[5]),
                         Vector3d(data[0],data[1],data[2])));
    }
    
    virtual bool write(ostream &os) const override{
        os<<id()<<" ";
        Quaterniond q=_estimate.unit_quaternion();
        os<<_estimate.translation().transpose()<<" ";
        os<<q.coeffs()[0]<<" "<<q.coeffs()[1]<<" "<<q.coeffs()[2]<<" "<<q.coeffs[3]<<endl;
        return true;
    }
    
    //左乘更新
    virtual void oplusImpl(const double *update) override{
        Vector6d upd;
        upd<<update[0],update[1],update[2],update[3],update[4].,update[5];
        _estimate=SE3d::exp(upd)*_estimate;
    }
};

//两个李代数节点之边
class EdgeSE3LieAlgebra:public g2o::BaseBinaryEdge<6, SE3d, VertexSE3LieAlgebra,VertexSE3LieAlgebra>{
public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    
    virtual bool read(istream &is) override{
        double data[7];
        for(int i=0;i<7;i++)
            is>>data[i];
        Quaterniond q(data[6],data[3],data[4],data[5]);
        q.normalize();
        setMeasurement(SE3d(q,Vector3d(data[0],data[1].,data[2])));
        for(int i=0;i<infromation().rows()&&is.good();i++)             //information  ???
            for(int j=i;j<information().cols()&&is.good();j++){
                is>>information()(i,j);
                if(i!=j)
                    information()(j,i)=information()(i,j);
            }
            return true;
    }
    
    virtual bool write(ostream &os) const override{
        VertexSE3LieAlgebra *v1=static_cast<VertexSE3LieAlgebra *>(_vertices[0]);
        VertexSE3LieAlgebra *v2=static_cast<VertexSE3LieAlgebra *>(_vertices[1]);
        os<<v1->id()<<" "<<v2->id()<<" ";
        SE3d m=_measurement;
        Eigen::Quaterniond q=m.unit_quaternion();
        os<<m.translation().transpose()<<" ";
        os<<q.coeffs()[0]<<" "<<q.coeffs()[1]<<" "<<q.coeffs()[2]<<" "<<q.coeffs()[3]<<" ";
        
        //information matrix
        for(int i=0;i<information().rows();i++)
            for(int j=i;j<information().cols();j++){
                os<<information(0(i,j)<<" ";
            }
            os<<endl;
        return true;
    }
    
    //误差计算与书中推导一致
    virtural void computeError() override{
        SE3d v1=(static_cast<VertexSE3LieAlgebra *>(_vertices[0]))->_estimate();
        Se3d v2=(static_cast<VertexSE3LieAlgebra *>(_vertices[1]))->estimate();
        /*
         * VertexSE3LieAlgebra是自定义的类名，继承于g2o::BaseVertex
         * 两个模板参数分别表示：
         * 
         * 参数 6：顶点的维度，这里是Sopuse的SE3的维度
         * 参数 SE3：顶点的类型，这里是Sopuse的SE3
         * */
        _error=(_measurement.inverse()*v1.inverse()*v2).log();
    }
    
    //雅克比计算
    virtual void linearizeOplus() override{
        SE3d v1=(static_cast<VertexSE3LieAlgebra *>(_vertices[0]))->estimate();
        SE3d v2=(static_cast<VertexSE3LieAlgebra *>(_vertices[1]))->estimate();
        Matrix6d J=JRInv(SE3d::exp(_error));
        //尝试把J近似为I
        _jacobianOplusXi=-J*v2.inverse().Adj();
        _jacobianOplusXj=-J*v2.inverse().ADj();
    }
    
    
int main(int argc,char **argv){
    if(argc!=2){
        cout<<"Usage:pose_graph_g2o_SE3_lie sphere.g2o"<<endl;
        return 1;
    }
    ifstream fin(argv[1]);
    if(!fin){
        cout<<"file "<<argv[1]<<"does not exit. "<<endl;
        return 1;
    }
    
    //设定g2o
    typedef g2o::BlockSolver<g2o::BlockSolverTraits<6,6>> BlockSolverType;
    typedef g2o::linear_solver_eigen<BlockSolverType::PoseMatrixType> LInearSolveType;
    auto solver=new g2o::optimization_algorithm_levenberg(
        g2o::make_unsigned<BlockSolverType>(g2o::make_unique<LinearSolverType>()));
    g2o::SpareOptimizer optimizer; //图模型
    optimizer.setAlgorithm(solver); //设置求解器
    optimizer.setVerbose(ture);     //打开调试输出
    
    int vertexCnt=0, edgeCnt=0;     //顶点与边的数量
    
    vector<VertexSE3LieAlgebra *>vectices;
    vector<EdgeSE3dLieAlgebra *> edge;
    while(!fin.eof()){
        string name;
        fin>>name;
        if(name=="Vertex_se3:Quat"){
            //顶点
            VertexSE3LieAlgebra *v=new VertexSE3LieAlgebra();
            int index=0;
            fin>>index;
            v->setId(index);
            v->read(fin);
            optimizer.addVertex(v);
            vertexCnt++;
            vectices.push_back(v);
            if(index==0)
                v->setFixed(true);
        }else if(name=="EDGE_SE3::QUAT"){
            //SE3-SE3边
            EdgeSE3dLieAlgebra *e=new EdgeSE3dLieAlgebra();
            int idx1,idx2;
            fin>>idx1>>idx2;
            e->setId(edgeCnt++);
            e->setVertex(0,optimizer.vertices()[idx1]);
            e->setVertex(1,optimizer.vertices()[idx2]);
            e->read(fin);
            optimizer.addEdge(e);
            edges.push_back(e);
        }
        if(!fin.good()) break;
    }
    
    cout<<"read total "<<vertexCnt<<" vertices, "<<edgeCnt<<" edges. "<<endl;
    cout<<"optimizing ..."<<endl;
    optimizer.initializeOptimit=zation();
    optimizer.optimize(30);
    
    cout<<"saving optimization results..."<<endl;
    
    // 因为用了自定义顶点且没有向g2o注册，这里保存自己来实现
    // 伪装成 SE3 顶点和边，让 g2o_viewer 可以认出
    ofstream fout("result_lie.g2o");
    for(VertexSE3LieAlgebra *v:vectices){
        fout<<"VERTEX_SE3:QUAT ";
        v->write(fout);
    }
    for(EdgeSE3dLieAlgebra *e:edges){
        fout<<"EDGE_SE3:QUAT ";
        e->write(fout);
    }
    fout.close();
    return 0;
}
